/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.fn.data;



import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.util.WindowedValue;

import com.bff.gaia.unified.vendor.grpc.v1p13p1.com.google.protobuf.ByteString;

import com.bff.gaia.unified.vendor.grpc.v1p13p1.io.grpc.stub.StreamObserver;

import com.bff.gaia.unified.vendor.guava.com.google.common.annotations.VisibleForTesting;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import java.io.IOException;



/**

 * A buffering outbound {@link FnDataReceiver} for the Unified Fn Data API.

 *

 * <p>Encodes individually consumed elements with the provided {@link Coder} producing a single

 * {@link UnifiedFnApi.Elements} message when the buffer threshold is surpassed.

 *

 * <p>The default buffer threshold can be overridden by specifying the experiment {@code

 * unified_fn_api_data_buffer_limit=<bytes>}

 *

 * <p>TODO: Handle outputting large elements (&gt; 2GiBs). Note that this also applies to the input

 * side as well.

 *

 * <p>TODO: Handle outputting elements that are zero bytes by outputting a single byte as a marker,

 * detect on the input side that no bytes were read and force reading a single byte.

 */

public class UnifiedFnDataBufferingOutboundObserver<T>

    implements CloseableFnDataReceiver<WindowedValue<T>> {

  // TODO: Consider moving this constant out of this class

  public static final String BEAM_FN_API_DATA_BUFFER_LIMIT = "unified_fn_api_data_buffer_limit=";

  @VisibleForTesting

  static final int DEFAULT_BUFFER_LIMIT_BYTES = 1_000_000;

  private static final Logger LOG =

      LoggerFactory.getLogger(UnifiedFnDataBufferingOutboundObserver.class);



  public static <T> UnifiedFnDataBufferingOutboundObserver<T> forLocation(

      LogicalEndpoint endpoint,

      Coder<WindowedValue<T>> coder,

      StreamObserver<UnifiedFnApi.Elements> outboundObserver) {

    return forLocationWithBufferLimit(

        DEFAULT_BUFFER_LIMIT_BYTES, endpoint, coder, outboundObserver);

  }



  public static <T> UnifiedFnDataBufferingOutboundObserver<T> forLocationWithBufferLimit(

      int bufferLimit,

      LogicalEndpoint endpoint,

      Coder<WindowedValue<T>> coder,

      StreamObserver<UnifiedFnApi.Elements> outboundObserver) {

    return new UnifiedFnDataBufferingOutboundObserver<>(

        bufferLimit, endpoint, coder, outboundObserver);

  }



  private long byteCounter;

  private long counter;

  private boolean closed;

  private final int bufferLimit;

  private final Coder<WindowedValue<T>> coder;

  private final LogicalEndpoint outputLocation;

  private final StreamObserver<UnifiedFnApi.Elements> outboundObserver;

  private final ByteString.Output bufferedElements;



  private UnifiedFnDataBufferingOutboundObserver(

      int bufferLimit,

      LogicalEndpoint outputLocation,

      Coder<WindowedValue<T>> coder,

      StreamObserver<UnifiedFnApi.Elements> outboundObserver) {

    this.bufferLimit = bufferLimit;

    this.outputLocation = outputLocation;

    this.coder = coder;

    this.outboundObserver = outboundObserver;

    this.bufferedElements = ByteString.newOutput();

    this.closed = false;

  }



  @Override

  public void close() throws Exception {

    if (closed) {

      throw new IllegalStateException("Already closed.");

    }

    closed = true;

    UnifiedFnApi.Elements.Builder elements = convertBufferForTransmission();

    // This will add an empty data block representing the end of stream.

    elements

        .addDataBuilder()

        .setInstructionReference(outputLocation.getInstructionId())

        .setTarget(outputLocation.getTarget());



    LOG.debug(

        "Closing stream for instruction {} and "

            + "target {} having transmitted {} values {} bytes",

        outputLocation.getInstructionId(),

        outputLocation.getTarget(),

        counter,

        byteCounter);

    outboundObserver.onNext(elements.build());

  }



  @Override

  public void flush() throws IOException {

    if (bufferedElements.size() > 0) {

      outboundObserver.onNext(convertBufferForTransmission().build());

    }

  }



  @Override

  public void accept(WindowedValue<T> t) throws IOException {

    if (closed) {

      throw new IllegalStateException("Already closed.");

    }

    coder.encode(t, bufferedElements);

    counter += 1;

    if (bufferedElements.size() >= bufferLimit) {

      flush();

    }

  }



  private UnifiedFnApi.Elements.Builder convertBufferForTransmission() {

    UnifiedFnApi.Elements.Builder elements = UnifiedFnApi.Elements.newBuilder();

    if (bufferedElements.size() == 0) {

      return elements;

    }



    elements

        .addDataBuilder()

        .setInstructionReference(outputLocation.getInstructionId())

        .setTarget(outputLocation.getTarget())

        .setData(bufferedElements.toByteString());



    byteCounter += bufferedElements.size();

    bufferedElements.reset();

    return elements;

  }

}